//
//  smoother.cpp
//  yuzon
//
//  Created by Alvin on 2023/3/29.
//

#include "smoother.hpp"
#include "utils.hpp"

using namespace std;

Smoother::Smoother() :
fc_min_(kFcMin),
b_(kB),
fc_d_(kFcD),
t_prev_(0),
x_hat_prev_(vector<float>(34, 0.0f)),
dx_hat_prev_(vector<float>(34, 0.0f)) {}

void Smoother::Process(KeypointResults &results) {
  if (results.empty())
    return;
  
  vector<float> &keypoints = results[0].keypoints;
  int64_t t = GetCurrentTime();
  
  for (int i = 0; i < results[0].num_joints; ++i) {
    if (t_prev_ == 0) {
      x_hat_prev_[i * 2] = keypoints[i * 3 + 1];
      x_hat_prev_[i * 2 + 1] = keypoints[i * 3 + 2];
    } else {
      keypoints[i * 3 + 1] = OneEuroFilter(t, keypoints[i * 3 + 1], i * 2);
      keypoints[i * 3 + 2] = OneEuroFilter(t, keypoints[i * 3 + 2], i * 2 + 1);
    }
  }
  
  t_prev_ = t;
}

void Smoother::Reset() {
  t_prev_ = 0;
  x_hat_prev_.clear();
  dx_hat_prev_.clear();
}

float Smoother::OneEuroFilter(int64_t t, float x, int i) {
  float t_e = (t - t_prev_) / 1000000.0f;
  
  float a_d = SmoothingFactor(t_e, fc_d_);
  float dx = (x - x_hat_prev_[i]) / t_e;
  float dx_hat = ExponentialSmoothing(a_d, dx, dx_hat_prev_[i]);
  
  float fc = fc_min_ + b_ * abs(dx_hat);
  float a = SmoothingFactor(t_e, fc);
  float x_hat = ExponentialSmoothing(a, x, x_hat_prev_[i]);
  
  x_hat_prev_[i] = x_hat;
  dx_hat_prev_[i] = dx_hat;
  
  return x_hat;
}

float Smoother::ExponentialSmoothing(float a, float x, float x_hat_prev) {
  return a * x + (1 - a) * x_hat_prev;
}

float Smoother::SmoothingFactor(float t_e, float fc) {
  float r = 2 * M_PI * fc * t_e;
  return r / (r + 1);
}
